Kimi K3 NPU Training

导言

Kimi K3 的 NPU 适配不是给现有 MLA-MoE 模型换一组配置。它同时引入 Kimi Delta Attention（KDA）、Block Attention Residuals（AttnRes）和 Stable LatentMoE，分别改变层内状态、跨层残差和专家通信。

截至 2026 年 7 月 20 日，官方已确认 K3 是 2.8T 参数、原生多模态、1M 上下文、896 专家激活 16 个，并采用 3× KDA + 1× Gated MLA；但完整权重、精确 config 和技术报告仍待发布。因此本文严格区分 已确认事实、组件证据、工程推导和发布后必验项，目标是形成可执行的 NPU bring-up 与性能优化计划，而不是制造一份猜测配置。

结论

K3 的 NPU 训练适配可以复用现有 MLA、MoE、GDN、FSDP2 和 checkpoint 基础设施，但不能按“K2 加一个新 attention”估算。真正决定工期的四个问题是：

1. KDA 训练算子：公开 FlashKDA 只有 CUDA 推理前向；NPU 需要 BF16 前向、反向、短卷积、递归状态、重计算和后续 TND/CP 语义。
2. AttnRes 跨层状态：它不再只传递相邻层 residual。PP 必须对齐 AttnRes block，并管理跨 stage 的增量 block representation。
3. 896 专家的分片：EP 是基础能力，不是可选优化。普通 FSDP2 整层 all-gather 无法合理承载 2.8T、896 专家模型。
4. 配置仍未公开：精确层数、hidden size、head、latent size、vision tower 和 checkpoint key 都是权重发布后的 release gate。

目标路径	建议	第一阻塞项	权重处理
Megatron + MindSpeed	作为完整 TP/PP/EP/CP 路径	KDA NPU、AttnRes PP、K3 model spec	HF 权重需要逻辑映射和分布式转换
FSDP2/HSDP + EP	作为更贴近 HF 的快速 bring-up 路径	EP-aware expert sharding、KDA NPU	可直接分片加载 HF；建议转 DCP 加速重启
定长数据	第一阶段默认	多模态 batch contract	工作量可控，利于静态 EP
TND/THD	第二阶段	KDA 序列边界、AttnRes cache、PP/CP metadata	不是只改一次 reshape

当前不能确定的参数

K3 的精确层数、hidden size、KDA/MLA head 数、latent size、vision tower、共享专家数、checkpoint key 和实际激活参数量尚未公开。Kimi Linear 48B-A3B 的 27 layers / hidden 2304 / 256 experts 只能用于理解代码形状，不能写入 K3 config。

{ width=95% }

自绘认知示意图：先把官方事实、组件证据和发布后必验项分开，再依次处理 KDA、AttnRes、LatentMoE 与 NPU 适配。

模型结构

已确认的骨架

官方材料确认 K3 使用三组相互耦合的新结构：[^k3-blog][^k3-quickstart]

• KDA 与 Gated MLA：按 3:1 周期交错。KDA 用固定状态压缩历史，Gated MLA 周期性提供全局 softmax attention。
• Block AttnRes：沿深度方向从多个 block representation 中学习加权读取，而不是固定把上一层 residual 与当前分支相加。
• Stable LatentMoE：共有 896 个 routed experts，每个 token 激活 16 个；专家主计算位于低维 latent space。
• 训练稳定性与效率组件：Quantile Balancing、Per-Head Muon、SiTU，以及从 SFT 阶段开始的 MXFP4 weight + MXFP8 activation QAT。

如果总 Transformer 层数是 L，忽略不完整尾部后只能得到近似层分布：

KDA layers       ≈ 3L / 4
Gated MLA layers ≈  L / 4

这不是精确层表。发布后必须从 layer_types 或等价配置逐层展开，并检查首尾层、vision 融合层和 AttnRes block 边界是否打破周期。

{ width=82% }

Kimi Linear 论文 Figure 3：KDA、MLA 与 MoE 的 3:1 混合结构。它是 K3 的组件先例，不是 K3 完整结构图；K3 还加入 Block AttnRes、Stable LatentMoE 等变化。

KDA 不是普通线性 attention 的改名。相较 Gated Delta Net（GDN）每个 head 使用标量 forget gate，KDA 为通道维度提供更细粒度的 gate。公开 Kimi Linear 实现的典型计算包含：[^kimi-linear][^fla-kda]

1. 从归一化输入投影 q/k/v、forget gate、update gate 和 output gate。
2. 对 q/k 分支执行短卷积，并对 packed sequence 边界做状态隔离。
3. 对 q/k 做 L2 normalization，构造通道级衰减与 beta。
4. 训练时执行 chunkwise KDA，decode 时更新每个 head 的固定矩阵状态。
5. 对 KDA 输出做 head-wise RMSNorm、output gate 和线性投影。

Kimi Linear 不是 K3 config

Kimi Linear 48B-A3B 有 27 层，其中 20 个 KDA 层和 7 个 full-attention 层；hidden size 为 2304，KDA 有 32 个 128 维 head，MoE 有 256 个 experts、top-8。本文只用它定位 KDA tensor 与 forward 分支。K3 已公开的是 2.8T、896/top-16 和 3:1 周期，不能把前者参数外推到后者。

参数量与内存

设非 routed-expert 参数为 D，routed-expert 参数为 E。官方总参数给出：

D + E = 2.8T

如果只按 16/896 的专家稀疏度计算，token 参与的参数近似为：

P_active = D + E × 16 / 896
         = 50B + 55D / 56

因此只能确定 routed expert 部分每次激活 1/56，不能确定 D。“K3 激活 50B”不是已公开事实；shared expert、attention、embedding、vision tower、latent projection 都会增加 active parameters。

从工程容量看，2.8T 已经给出硬下界：

• BF16 权重：约 5.6 TB，还不含 padding、metadata 与临时转换空间。
• 常见 mixed-precision Adam 规划代理：按约 16 bytes/parameter 是 44.8 TB 全局状态；K3 实际使用 Per-Head Muon，精确状态必须按公开实现重算。
• 分摊下界：44.8 TB 在 512、1024、2048 卡上分别约 87.5、43.75、21.9 GB/卡，尚未计 activation、通信 buffer、FSDP all-gather 和 NPU workspace。

这说明首个正确性任务必须使用 tiny proxy config。不应把完整 2.8T checkpoint 当作算子单测环境。

前向流程

在未知精确 vision contract 的前提下，可以建立如下模型级前向假设；虚线项都需要发布后核验：

flowchart LR
  A["文本与视觉输入"] --> B["Tokenizer / Vision Tower / Projector"]
  B --> C["Token Embedding"]
  C --> D["Block AttnRes 深度读取"]
  D --> E{"当前层类型"}
  E -->|"3/4 周期"| F["KDA: Conv + Gate + Delta State"]
  E -->|"1/4 周期"| G["Gated MLA"]
  F --> H["Stable LatentMoE"]
  G --> H
  H --> I["更新 AttnRes block representation"]
  I --> D
  I --> J["Final Norm / LM Head"]

三类状态决定并行实现：

• 序列状态：KDA 的短卷积缓存和矩阵递归状态，需要正确处理 sequence boundary、CP 切分和重计算。
• 深度状态：Block AttnRes 的 block representation 跨层存在，PP 不再只传一个 hidden tensor。
• 专家状态：Stable LatentMoE 的 router、latent projection、expert bucket 和负载统计跨 EP group 组织。

NPU 算子支持度

现有 Ascend 栈并非从零开始。RMSNorm、RoPE、softmax attention、grouped matmul、MoE token permute/unpermute 和 MC2 路径都已有基础；MindSpeed-LLM 还提供 Qwen3-Next GDN 的 NPU、CP、packed/TND 与 FSDP2+EP 参考实现。[^mindspeed-gdn][^mindspeed-fsdp-ep]

模块	可复用基础	当前支持判断	主要缺口
RMSNorm / linear / SiLU	torch_npu 与融合 MLP	高	K3 精确 fusion 边界
Gated MLA	NPU fused attention、RoPE、K2 MLA 经验	中高	Gated MLA 的 gate 与投影布局
LatentMoE	GMM、permute/unpermute、AllToAll/MC2	中	latent projection、SiTU、router 与 shared expert 语义
Block AttnRes	基础 GEMM/softmax/RMSNorm	中低	跨层 cache、PP 通信与融合算子
KDA 固定长度前向	GDN/FLA 算法和 NPU GDN 先例	低	通道 gate、DPLR chunk kernel、短卷积、状态布局
KDA 反向	FLA CUDA autograd 语义	缺失	dq/dk/dv/dbeta/dgate/dstate/dA/dbias 的 NPU kernel 与精度
KDA TND/CP	GDN CP、FLA cu_seqlens 参考	缺失	packed 边界 reset、head/sequence 通信、变长 backward
MXFP4/MXFP8 训练	官方 K3 QAT 事实、NPU 低精度基础	待核验	training recipe、scale granularity、checkpoint 与累计精度

不要把 FlashKDA 当训练支持

FlashKDA 当前面向 NVIDIA SM90+、CUDA 12.9+，公开路径在 torch.inference_mode() 下只做 forward。它支持 fixed/varlen 和 BF16，可以说明算子接口与 inference 性能潜力，但不能证明 NPU、backward、recompute 或训练稳定性。[^flash-kda]

Kimi Linear 的公开组件结果表明混合 KDA 在 1M 上下文可减少最多 75% KV cache，并在论文环境中获得最高约 6.3 倍 decode throughput。它解释了 K3 选择 KDA 的动机，不是 K3 在 Ascend 上的性能承诺。[^kimi-linear]

{ width=82% }

Kimi Linear 论文 Figure 7：48B 组件模型在论文硬件与实现下的 prefill 和 decode 对比。该证据用于确定 KDA 优化价值，不可直接外推为 K3 NPU 加速比。

Megatron 后端

权重是否需要转换

如果从官方 HF checkpoint 开始，答案是：需要逻辑权重转换，但不应做单机集中式转存。

• HF 与 Megatron 对 QKV/MLA、experts、TP shard、PP layer placement 和 distributed checkpoint 的 tensor 命名与布局不同。
• 2.8T BF16 权重已经约 5.6 TB，集中加载再 torch.cat/split 会制造不可接受的主机内存和临时磁盘峰值。
• 应扩展 Megatron Bridge，按 tensor 流式读取 safetensors、执行必要的 transpose/concat/split，并直接写入 distributed checkpoint。
• 转换必须提供逐 tensor shape/dtype/hash 或数值抽样，以及 HF → Megatron → HF 的 round-trip 验证。

如果是从头预训练，不存在“HF 权重转 Megatron”的起始步骤；但仍应保存 Megatron distributed checkpoint。已有 Kimi K2 Bridge 可以复用 MLA/MoE 映射模式，K3 仍需新增 KDA、AttnRes、Stable LatentMoE、vision tower 和精度 metadata。[^megatron-bridge]

模型与并行改造

Megatron 当前已有 GDN 的 TP、CP、packed THD、selective recompute 和 distributed checkpoint 支持，也已有 moe_latent_size 的 LatentMoE 测试。它们是 K3 的骨架，不是现成 K3 支持。[^megatron-gdn][^megatron-latent]

能力	设计建议	关键验证
TP	KDA 按 head 分片 q/k/v/gate/state；Gated MLA 按 head/latent projection 分片	TP=1/2/4 的 loss、梯度和 recurrent state parity
PP	stage 边界对齐 4 层 attention 周期和 AttnRes block	block cache 所有权、跨 stage 增量通信、virtual PP 与 checkpoint
EP	896 experts 按拓扑分组，优先使用 GMM + permute/unpermute + MC2	expert placement、token drop、capacity、负载、跨节点 A2A
CP	MVP 设为 1；后续参考 GDN head-parallel all-to-all 或 FLA CP	sequence chunk、conv/state continuity、Gated MLA 与 KDA 共存
重计算	selective recompute；KDA projection/局部状态按收益重算	router 决策和 AttnRes cache 必须确定；避免重放主导 A2A

896 的因数对 EP 规划很友好，但仍不能只按整除关系选：

EP=64  -> 14 experts/rank
EP=112 ->  8 experts/rank
EP=128 ->  7 experts/rank
EP=224 ->  4 experts/rank

最终值取决于 NPU supernode 规模、机内/机间带宽、expert TP、micro batch、top-16 token expansion 和 latent size。官方推荐 64 卡以上 supernode，并强调 fully balanced EP 使用静态 shape、关键路径无 host sync；这意味着 静态 expert bucket 与通信拓扑 应进入第一版设计，而不是性能阶段再补。[^k3-blog]

AttnRes 对 PP 的影响

普通 pre-norm Transformer 的 PP stage 只需传当前 hidden state。Block AttnRes 会读取多个历史 block representation；若直接把全部历史跨 stage 传递，通信和存储随深度增长。

Attention Residuals 论文给出的方向是：

1. 在 block 内维持局部 residual。
2. 只在 block 完成时产生稳定 representation。
3. PP stage 增量缓存已完成 block，而不是每层重发全历史。
4. layer query 对缓存 representation 做归一化、softmax 和加权读取。

因此 PP 切分必须优先落在 AttnRes block 边界，并尽量同时对齐 KDA-KDA-KDA-MLA 周期。论文中的 <4% PP 开销来自组件实验，不能当 K3 集群验收阈值。[^attnres]

FSDP2 后端

FSDP2 的优势是更接近 HF 模型代码和 state dict。可以在 meta device 上实例化模型，再按 mesh 分片加载 safetensors；因此 不强制先生成 Megatron 格式。为了故障恢复和大规模重复启动，仍建议一次性写成 PyTorch DCP 或等价 sharded checkpoint。

真正的问题是专家层：

• 若把整个 MoE block 当普通 FSDP unit，forward 前可能 all-gather 896 experts，失去稀疏模型的内存意义。
• 若按动态路由到的 expert 独立触发 FSDP collectives，不同 rank 的调用集合和顺序可能不一致，导致 collective mismatch 或串行抖动。
• 可行结构应是 dense/attention 参数走 FSDP2 或 HSDP，routed experts 走静态 EP placement 与专用 dispatcher；expert 内是否再 fully shard 取决于单 expert 尺寸和节点层级。

当前 MindSpeed-LLM 的 FSDP2 路径已经包含 EP dispatcher、GMM/MC2 和 GDN context-parallel 组件，是比“原生 PyTorch FSDP 包住 HF K3”更接近目标的底座。[^mindspeed-fsdp-ep]

FSDP2 工作项	说明	规模
HF 模型接入与 meta init	解析 K3 config、空权重构造、auto-wrap policy	M
Sharded load / DCP	safetensors 分片加载、保存恢复、精度 metadata	M–L
Dense HSDP/FSDP2	attention、embedding、vision 和 shared parameter 分片	M
EP-aware experts	静态 expert placement、dispatcher、GMM/MC2、shared expert	L–XL
KDA NPU	与 Megatron 共用核心算子，但 adapter/state 管理不同	XL
AttnRes	FSDP wrapping、checkpoint 与跨层 cache	L
TND/CP	packed metadata 与 KDA state 协同	XL

FSDP2 路径减少 model spec 和 HF 转换工作，不会消除核心算子工作；其最大价值是更快建立 单机/小规模 reference 与固定长度闭环。

定长与 TND

先做定长

第一阶段建议：BF16、定长 padding、静态 shape、CP=1、固定 expert bucket。

原因不是定长最终性能最好，而是它把变量压缩到最少：

• KDA 短卷积和 recurrent state 不需要在一个 tensor 内处理多个序列边界。
• Gated MLA 可以直接调用稳定的 fused attention 接口。
• AttnRes cache 与 PP metadata 不需要跟随 packed segment 改变。
• MoE 可以提前分配静态 token bucket，更接近官方 fully balanced EP 的描述。
• BF16 先建立收敛与梯度 reference，避免把 KDA、并行 bug 和 MXFP4/MXFP8 量化误差混在一起。

TND 不是布局开关

TND/THD packed 训练至少新增六组工作：

1. 数据：cu_seqlens、position、vision span、label 和 loss mask 必须同序重排。
2. KDA：short conv 和 recurrent state 在每个 packed sequence 边界 reset；forward、backward 和 recompute 一致。
3. Gated MLA：接入 varlen fused attention，并验证多模态 span。
4. AttnRes：block cache 采用 token-major 表示，不能跨 packed sample 泄漏。
5. MoE：动态有效 token 数与静态 expert bucket、Quantile Balancing 同时成立。
6. 分布式：PP/CP 的 metadata、checkpoint 和 loss normalization 保留 sequence boundary。

公开 FlashKDA 的 varlen 是 inference forward；FLA 提供 CUDA training autograd 与 cu_seqlens，而 MindSpeed GDN 提供 NPU CP/TND 工程先例。三者合起来说明路径可行，但没有一个直接等于 K3 KDA NPU TND training。

开箱性能优化

性能工作应按依赖顺序开展，而不是一开始同时打开 TP/PP/EP/CP、TND 和低精度。

优先级

1. P0 正确性
   • tiny config 完成 KDA BF16 forward/backward、state、recompute parity。
   • 发布后冻结 K3 config、逐层类型、tensor map、vision contract 和 active parameters。
   • 单 rank 跑通固定长度 loss 与 optimizer step。
2. P1 并行闭环
   • EP 拓扑对齐，并优先接 NPU GMM、permute/unpermute、AllToAll/MC2。
   • TP 分片 KDA/MLA；PP 同时对齐 attention 周期和 AttnRes block。
   • CP 暂时保持 1。
3. P2 内存与调度
   • selective recompute、distributed checkpoint、FSDP async prefetch。
   • overlap shared/routed expert、token dispatch 与 expert compute。
   • 避免在重计算阶段重复主导通信。
4. P3 算子融合
   • 首要优化 KDA chunk forward/backward、短卷积与 gate 融合。
   • 其次是 LatentMoE projection + SiTU + GMM、Gated MLA 与 AttnRes depth attention。
5. P4 长序列
   • 引入 TND，再实现 KDA-aware CP，测量 padding 节省是否覆盖新增通信和动态 shape 成本。
6. P5 低精度
   • 只有 BF16 loss/gradient/收敛基线稳定后，再启用 MXFP4/MXFP8 QAT 或官方训练 checkpoint 路径。

第一组性能实验

用 tiny/small K3 proxy 固定 TP=1, PP=1, CP=1，只扫描 EP、micro batch、sequence length 和 expert bucket；再分别加入 TP、PP。这样可以从 profiler 中分离 KDA kernel、MoE AllToAll/GMM 和 AttnRes 通信，避免一次实验出现四个并行轴的混合瓶颈。

建议观测指标

• KDA forward/backward 的 kernel time、带宽、workspace、chunk size 与状态读写量。
• Gated MLA 与 KDA 层的逐层耗时比例，验证 3:1 层分布是否形成预期瓶颈。
• MoE dispatch/combine、A2A、GMM、shared expert overlap、每 expert token p50/p99。
• AttnRes depth-softmax、cache 内存、跨 PP stage bytes 与 bubble。
• MFU、tokens/s、有效 tokens/s、padding ratio、recompute FLOPs 和峰值 HBM。
• BF16 与低精度下的 loss、gradient norm、router entropy、expert imbalance 和溢出/下溢计数。

需求拆解

规模定义：S <= 1 人周，M = 1–3 人周，L = 3–6 人周，XL > 6 人周。它假设人员熟悉 Megatron/MindSpeed/NPU 算子，并且不包含集群排队和上游发布延期。

Epic	交付物	规模	依赖
发布物审计	exact config、层表、参数计算、dtype、key inventory	M	K3 权重与报告
Reference 模型	HF/CPU 或 GPU reference、tiny proxy、确定性 loss	L	config 审计
KDA NPU	定长 BF16 forward/backward、state、recompute	XL	KDA reference
Block AttnRes	单 stage、block cache、深度 attention 测试	L	exact block config
Stable LatentMoE	latent projection、router、SiTU、shared/routed expert	L	exact MoE config
Gated MLA	gate、projection、fused attention	M	exact MLA config
多模态	vision tower/projector、placeholder/span、loss mask	L，待定	公开代码
Megatron Bridge	分布式流式转换、round-trip 校验	L	exact checkpoint keys
TP	KDA/MLA/LatentMoE tensor shard 与 parity	M	model + kernels
EP	expert placement、dispatcher、GMM/MC2、负载指标	L	Stable LatentMoE
PP + AttnRes	block-aligned schedule、跨 stage cache	L–XL	AttnRes
CP	KDA-aware sequence/head parallel	XL	定长 KDA parity
重计算	deterministic router/cache、selective policy	L	完整 forward/backward
FSDP2/HSDP + EP	dense sharding、expert-aware sharding、DCP	L–XL	HF reference
定长数据	文本/多模态静态 batch、loss normalization	M	tokenizer/vision contract
TND/THD	全链路 packed boundary 与并行 metadata	XL	定长闭环
精度与性能	BF16 基线、低精度、fusion、规模化调优	持续	以上全部

阶段验收

1. M0：发布审计
   • 精确 config、参数量、层分布和 tensor key 可重复计算。
   • 区分 BF16、MXFP4/MXFP8 和推理/训练 checkpoint。
2. M1：单 rank 定长
   • tiny config 的 logits、loss、关键中间 tensor 与 reference 对齐。
   • KDA 全量梯度、AttnRes cache 和 MoE routing 测试通过。
3. M2：TP/EP
   • TP=1/2/4、EP=1/多 rank 的 loss 和 gradient parity 通过。
   • 无整层 896 expert all-gather，expert placement 可检查。
4. M3：PP 与规模化
   • block-aligned PP 恢复、重计算、distributed checkpoint 通过。
   • profiler 能分别归因 KDA、MoE 与 AttnRes 瓶颈。
5. M4：TND/CP/低精度
   • packed sequence 无状态串扰，CP parity 通过。
   • MXFP4/MXFP8 相对 BF16 的收敛与吞吐收益有端到端数据。

按当前未知量，固定长度 MVP 可暂按 2–3 名核心工程师、8–12 周估算；TND/CP 和成熟开箱性能再增加 6–10 周。这是立项代理值，不是承诺：KDA NPU backward、AttnRes PP 细节和官方 checkpoint 形态都可能改变关键路径。

总结

K3 的模型级特点可以概括为三种稀疏或压缩：KDA 压缩序列历史、AttnRes 选择性读取深度历史、LatentMoE 稀疏并压缩专家计算。它们同时把系统复杂度从“矩阵乘与 all-reduce”推向 状态、路由和跨层缓存。

因此最稳妥的 NPU 路径是：先冻结发布事实，再以 tiny config 建立 BF16 定长 reference；优先完成 KDA backward 与 EP-aware MoE；随后接 TP 和 block-aligned PP；最后处理 TND、KDA-aware CP 和 MXFP4/MXFP8。Megatron 与 FSDP2 可以共享核心 NPU 算子，但分别承担分布式转换/PP 和 expert-aware sharding 的主要成本。

[^k3-blog]: Kimi K3 technical blog
[^k3-quickstart]: Kimi K3 quick start
[^kimi-linear]: Kimi Linear technical report and official repository
[^fla-kda]: Flash Linear Attention KDA training implementation
[^flash-kda]: FlashKDA repository
[^attnres]: Attention Residuals and official repository
[^mindspeed-gdn]: MindSpeed-LLM GDN context parallel
[^mindspeed-fsdp-ep]: MindSpeed-LLM FSDP2 expert parallel
[^megatron-bridge]: Megatron Bridge Kimi support
[^megatron-gdn]: Megatron Core Gated Delta Net
[^megatron-latent]: Megatron Core LatentMoE unit test